import pandas as pd
import matplotlib.pyplot as plt
import numpy as np

FILE_PATH = "/home/shaonian/SED/SED/configs/dataset/desed_tsv/train_unlabeled_id_Y800M_thd95_200h.tsv"

df = pd.read_csv(FILE_PATH, sep="\t")
print(len(df))
# weak
events = df["event_labels"].apply(lambda x: x.split(",") if pd.notna(x) else [])
events = events.explode("event_labels")
print(df.head())
### strong
# events = df.groupby("filename")["event_label"].apply(lambda x: np.unique(x.dropna().tolist()))
# events = events.reset_index()
# check events that only have one label
# events = events[events["event_label"].apply(lambda x: len(x) < 2)]
# print(len(events), events)
# reconvert to filename-event_label pairs
# events = events.explode("event_labels")
###
# convert into dataframe
label_counts = events.value_counts(normalize=False).to_dict()
print(label_counts)
plt.title("Number of labels in non-single speech labels")
plt.xlabel("Labels")
plt.ylabel("Count") 
plt.figure(figsize=(10, 6))
# Convert weak_label_counts to percentages

total = sum(label_counts.values())
label_percentages = {k: (v/total)*100 for k, v in label_counts.items()}
weak_label_counts = {
    "speech": 550,
    "dog": 214,
    "cat": 173,
    "alarm_bell_ringing": 205,
    "dishes": 184,
    "frying": 171,
    "blender": 134,
    "running_water": 343,
    "vacuum_cleaner": 167,
    "electric_shaver_toothbrush": 103
}

weak_label_counts = {k: v for k, v in weak_label_counts.items()}
weak_total = sum(weak_label_counts.values())
weak_label_percentages = {k: (v/weak_total)*100 for k, v in weak_label_counts.items()}
# Convert labels_count to percentages
labels_percentages = pd.Series(label_percentages)
labels_percentages = labels_percentages.rename(lambda x: x.lower())

# Get all unique labels for consistent x-axis
all_labels = sorted(list(set(list(weak_label_percentages.keys()) + list(labels_percentages.index))))

# Prepare data for bar plot
weak_values = [weak_label_percentages.get(label, 0) for label in all_labels]
desed_values = [labels_percentages.get(label, 0) for label in all_labels]

x = range(len(all_labels))
width = 0.35

plt.bar([i - width/2 for i in x], weak_values, width, label='Weak labels', alpha=0.7)
plt.bar([i + width/2 for i in x], desed_values, width, label='Y800M labels', alpha=0.7)
plt.xticks(x, all_labels)
plt.title("Number of labels in non-single speech labels")
plt.xlabel("Labels")
plt.ylabel("Percentage (%)") 
plt.xticks(rotation=90)
plt.tight_layout()
plt.legend()
plt.savefig("./label_histogram.png")